import os
import shutil
import random
from pathlib import Path

def create_directories(base_path):
    """创建train、test、valid目录"""
    for dir_name in ['train', 'test', 'valid']:
        for subdir in ['images', 'labels']:
            path = os.path.join(base_path, dir_name, subdir)
            os.makedirs(path, exist_ok=True)

def split_dataset(source_dir, dest_dir, img_extensions=('.jpg', '.jpeg', '.png')):
    """
    将数据集分割为训练集、测试集和验证集
    
    Args:
        source_dir: 源数据目录，包含图片和标注文件
        dest_dir: 目标根目录，将在其下创建train/test/valid子目录
        img_extensions: 图片文件扩展名元组
    """
    # 创建目标目录
    create_directories(dest_dir)
    
    # 获取所有图片文件
    image_files = []
    for ext in img_extensions:
        image_files.extend(list(Path(source_dir).glob(f'*{ext}')))
    
    # 随机打乱文件列表
    random.shuffle(image_files)
    
    # 计算分割点
    total = len(image_files)
    train_split = int(0.8 * total)
    test_split = int(0.9 * total)
    
    # 分配数据集
    splits = {
        'train': image_files[:train_split],
        'test': image_files[train_split:test_split],
        'valid': image_files[test_split:]
    }
    
    # 移动文件
    for split_name, files in splits.items():
        print(f'Moving files to {split_name} directory...')
        for img_path in files:
            # 处理图片文件
            img_filename = img_path.name
            new_img_path = os.path.join(dest_dir, split_name, 'images', img_filename)
            shutil.copy2(str(img_path), new_img_path)
            
            # 处理对应的标注文件
            txt_filename = img_path.stem + '.txt'
            txt_path = img_path.parent / txt_filename
            if txt_path.exists():
                new_txt_path = os.path.join(dest_dir, split_name, 'labels', txt_filename)
                shutil.copy2(str(txt_path), new_txt_path)
            else:
                print(f'Warning: No annotation file found for {img_filename}')
        
        print(f'{split_name}: {len(files)} images')

if __name__ == '__main__':
    # 设置源数据目录和目标目录
    source_directory = "J:/皮尺点"#input("请输入源数据目录路径: ")
    destination_directory = "E:/code/fork/yolov10/mydata/pichidian2015.v1i.yolov8"#input("请输入目标目录路径: ")
    
    # 执行数据集分割
    split_dataset(source_directory, destination_directory)
    print("数据集分割完成！") 